Imports System.Data.Odbc
Imports System.IO
Imports System.Security.Cryptography

Public Class Database
    Private Shared connDriver As String = "{MYSQL ODBC 3.51 Driver}"
    Private Shared connServer As String = "physcomp1.wlu.ca"
    Private Shared connDatabase As String = "rusu4904"
    Private Shared connUsername As String = "rusu4904"
    Private Shared connPassword As String = "armc53fry"
    Private Shared connOption As String = "3"
    Private Shared conn As New System.Data.Odbc.OdbcConnection( _
        "DRIVER=" & connDriver & ";SERVER=" & connServer & _
        ";DATABASE=" & connDatabase & ";USER=" & connUsername & _
        ";PASSWORD=" & connPassword & ";OPTION=" & connOption)

    Public Shared db = New Database
    Private Shared myId As Integer = -1
    Private Shared myEmail As String = ""
    Private Shared myPrivateKey() As Byte = Nothing
    Public progress As ProgressBar

    Public Function testConnection() As Boolean
        Try
            If conn.State = ConnectionState.Closed Then
                conn.Open()
                Return True
            End If
        Catch e As Exception
        End Try
        Return False
    End Function

    Public Function login(ByVal email As String, ByVal password As String) As Boolean
        Try
            If email.Length < 5 Then Return False
            If password.Length < 3 Then Return False
            If conn.State = ConnectionState.Closed Then conn.Open()
            Dim cmd As New OdbcCommand("SELECT id, password, salt FROM ss_user WHERE email='" & email & "'", conn)
            Dim reader As OdbcDataReader = cmd.ExecuteReader

            reader.Read()
            Dim salt() As Byte = Convert.FromBase64String(reader.GetString(2))
            If Encryption.hashPassword(password, salt) = reader.GetString(1) Then
                myId = reader.GetInt32(0)
                myEmail = email
                Return True
            End If
        Catch e As Exception
            MsgBox("login: " & e.Message)
        End Try
        Return False
    End Function

    Public Function getMyId() As Integer
        Return myId
    End Function

    Public Function getMyEmail() As String
        Return myEmail
    End Function

    Public Function getMyPrivateKey() As Byte()
        Return myPrivateKey
    End Function

    Public Sub setMyPrivateKey(ByVal newKey() As Byte)
        myPrivateKey = newKey
    End Sub

    Public Function validateMyPrivateKey() As Boolean
        If myId = -1 Then Return False
        If myPrivateKey Is Nothing Then Return False
        Return Encryption.rsaCheckKeys(getPublicKey(myId), myPrivateKey)
    End Function

    Public Function getUsers() As User()
        If conn.State = ConnectionState.Closed Then conn.Open()
        Dim cmd As New OdbcCommand("SELECT * FROM ss_user where id !=" & myId, conn)
        Dim reader As OdbcDataReader = cmd.ExecuteReader

        Dim uArray(100) As User
        Dim count As Integer = 0
        Dim tempUser As User
        While reader.Read()
            tempUser = New User
            tempUser.id = CInt(reader("id"))
            tempUser.email = CStr(reader("email"))
            uArray(count) = tempUser
            count += 1
        End While

        'conn.Close()
        ReDim Preserve uArray(count - 1)
        Return uArray
    End Function

    Public Function checkForFile() As Boolean
        If conn.State = ConnectionState.Closed Then conn.Open()
        Dim cmd As New OdbcCommand("SELECT filename FROM ss_file where toUserId = '" & myId & "'", conn)
        Dim reader As OdbcDataReader = cmd.ExecuteReader
        Return reader.HasRows
    End Function

    Public Function checkForEmail(ByVal email As String) As Boolean
        If conn.State = ConnectionState.Closed Then conn.Open()
        Dim cmd As New OdbcCommand("SELECT email FROM ss_user where email = '" & email.Trim & "'", conn)
        Dim reader As OdbcDataReader = cmd.ExecuteReader
        Return reader.HasRows
    End Function

    Public Function getFileList() As String()
        If conn.State = ConnectionState.Closed Then conn.Open()
        Dim cmd As New OdbcCommand("SELECT filename FROM ss_file where toUserId = '" & myId & "'", conn)
        Dim reader As OdbcDataReader = cmd.ExecuteReader

        Dim fArray(100) As String
        Dim count As Integer = 1
        Dim tempFileName As String
        While reader.Read()
            tempFileName = CStr(reader("filename"))
            fArray(count) = tempFileName
            count += 1
        End While

        ReDim Preserve fArray(count - 1)
        Return fArray
    End Function

    Public Function createAccount(ByVal email As String, ByVal password As String) As Boolean
        Try
            If email.Length < 5 Then Return False
            If password.Length < 3 Then Return False
            Dim publicKey() As Byte = Nothing
            Dim privateKey() As Byte = Nothing
            myEmail = email
            Encryption.generateRsaKeys(publicKey, privateKey)
            myPrivateKey = privateKey
            Dim salt() As Byte = Encryption.generateSalt
            If conn.State = ConnectionState.Closed Then conn.Open()
            Dim cmd As New OdbcCommand()
            cmd.Connection = conn

            cmd.CommandText = "INSERT INTO ss_user (email, password, salt, publicKey) VALUES('" _
               & email & "', '" & Encryption.hashPassword(password, salt) & "', '" _
               & Convert.ToBase64String(salt) & "', '" & Convert.ToBase64String(publicKey) & "')"

            cmd.ExecuteNonQuery()
            'cmd.Dispose()
            'cmd = New OdbcCommand("SELECT id FROM ss_user WHERE email = '" & email & "'", conn)
            'Dim reader As OdbcDataReader = cmd.ExecuteReader
            'myId = reader.GetInt32(0)

            Return optimizeTable("ss_user")
        Catch e As Exception
            MsgBox("createAccount: " & e.Message)
        End Try
        Return False
    End Function

    Public Function optimizeTable(ByVal tableName As String) As Boolean
        Try
            If conn.State = ConnectionState.Closed Then conn.Open()
            Dim cmd As New OdbcCommand("OPTIMIZE TABLE `" & tableName & "`", conn)
            cmd.ExecuteNonQuery()
            Return True
        Catch e As Exception
            MsgBox("optimizeTable: " & e.Message)
        End Try
        Return False
    End Function

    Public Function getPublicKey(ByVal accountId As Integer) As Byte()
        Try
            If conn.State = ConnectionState.Closed Then conn.Open()
            Dim cmd As New OdbcCommand("SELECT publicKey FROM ss_user WHERE id = '" & accountId & "'", conn)
            Dim reader As OdbcDataReader = cmd.ExecuteReader
            reader.Read()
            Return Convert.FromBase64String(reader.GetString(0))
        Catch e As Exception
            MsgBox("getPublicKey: " & e.Message)
        End Try
        Return Nothing
    End Function

    Public Function uploadFile(ByVal filePath() As String, ByVal receiverId As Integer) As Boolean
        For i As Integer = 0 To filePath.Length - 1
            If Not uploadFile(filePath(i), receiverId) Then Return False
        Next
        Return True
    End Function

    Public Function uploadFile(ByVal filePath As String, ByVal receiverId As Integer) As Boolean
        Try
            If Not IO.File.Exists(filePath) Then Return False
            Dim tempFile As String = Path.GetTempFileName()
            Dim aesKey As Byte() = Nothing
            If Not Encryption.aesCrypt(filePath, tempFile, aesKey, Encryption.CryptFunction.encrypt) Then Return False

            Dim publicKey() As Byte = getPublicKey(receiverId)
            Dim encryptedKey() As Byte = Encryption.rsaEncrypt(aesKey, publicKey)
            Dim keyAsString As String = Convert.ToBase64String(encryptedKey)

            Dim tsString As String = Now.Ticks.ToString
            Transfer.uploadFile(tempFile, tsString, progress)
            Dim fileName As String = IO.Path.GetFileName(filePath)
            IO.File.Delete(tempFile)

            If conn.State = ConnectionState.Closed Then conn.Open()
            Dim cmd As New OdbcCommand()
            cmd.Connection = conn
            cmd.CommandText = "INSERT INTO ss_file(fromUserId, toUserId, serverFile, fileName, theKey) VALUES (" _
                & "'" & myId & "', '" & receiverId & "', '" & tsString & "', '" & fileName & "','" & keyAsString & "')"
            cmd.ExecuteNonQuery()
            Return optimizeTable("ss_file")
        Catch e As Exception
            MsgBox("uploadFile: " & e.Message)
        End Try
        Return False
    End Function

    Public Function downloadFile(ByVal serverFile As String, ByVal localFile As String) As Boolean
        Try
            If conn.State = ConnectionState.Closed Then conn.Open()
            Dim cmd As New OdbcCommand("SELECT theKey, serverFile FROM ss_file WHERE fileName = '" & serverFile & "'", conn)
            Dim reader As OdbcDataReader = cmd.ExecuteReader
            reader.Read()
            Dim key() As Byte = Encryption.rsaDecrypt(Convert.FromBase64String(reader.GetString(0)), myPrivateKey)

            Dim tempFile As String = Path.GetTempFileName()
            If Not Transfer.downloadFile(reader.GetString(1), tempFile, progress) Then Return False
            If Not Transfer.deleteFile(reader.GetString(1)) Then Return False

            If Not Encryption.aesCrypt(tempFile, localFile, key, Encryption.CryptFunction.decrypt) Then
                MsgBox("decryption failed")
                Return False
            End If

            IO.File.Delete(tempFile)
            cmd.Dispose()

            cmd = New OdbcCommand("DELETE FROM ss_file WHERE filename = '" & serverFile & "'", conn)
            cmd.ExecuteNonQuery()
            cmd.Dispose()
            Return True
        Catch e As Exception
            MsgBox("downloadFile: " & e.Message)
        End Try
        Return False
    End Function
End Class